#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import math
# from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from diff_gauss import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from utils.sh_utils import eval_sh

def render(viewpoint_camera, 
           pc: GaussianModel, 
           pipe, 
           bg_color: torch.Tensor, 
           theta: torch.Tensor,
           phi: torch.Tensor,
           density: torch.Tensor,
           scaling_modifier=1.0, 
           override_color = None):
    """
        Render the scene. 
        Background tensor (bg_color) must be on GPU!
    """
 
    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug
    )
        
    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    visible_mask = rasterizer.markVisible(positions=pc.get_xyz)
    sun_color = pc.feature_encode(viewpoint_camera, theta, phi, density, visible_mask)
    
    means3D = pc.get_xyz[visible_mask]
    means2D = screenspace_points[visible_mask]
    opacity = pc.get_opacity[visible_mask]
    rotations = pc.get_rotation[visible_mask]
    scales = pc.get_scaling[visible_mask]
    

    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    
    rendered_image, _, _, _, radii, sunpos_map = rasterizer(
                                        means3D = means3D,
                                        means2D = means2D,
                                        shs = None,
                                        colors_precomp = torch.ones_like(scales),
                                        opacities = opacity,
                                        scales = scales,
                                        rotations = rotations,
                                        cov3Ds_precomp = None,
                                        extra_attrs=sun_color)
    
    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "sunpos_map": sunpos_map,
            "viewspace_points": screenspace_points,
            "visibility_filter" : radii > 0,
            "radii": radii}
